# @Author: Ruian Duan 
# @Date: 2014/04/18
#
# Give a set of malicious classes, GroundTruth simply pick from each class COUNT samples.
# During selection, it will do random selection and filter out urls that land in same
# page or failed to connect.

import commands
import random
import utils
import os


class GroundTruth:
	def __init__(self, basedir = '/space/webcrawl-malware/data/TempData/LOGs/',
			jslistfile = 'mjslist.log',
			jshashfile = 'mjshashs',
			jshash2countfile = 'mjshash2count'):
		self.basedir = basedir
		self.jslistfile = self.basedir + jslistfile
		self.jshashfile = self.basedir + jshashfile
		self.jshash2countfile = self.basedir + jshash2countfile
		self.jslist = utils.DictFile().file2dict(self.jslistfile, 'None').keys()
		self.jslist = [item.split('/')[-1] for item in self.jslist]
		self.jshashs = utils.DictFile().file2dict(self.jshashfile, 'str')
		self.jshash2count = utils.DictFile().file2dict(self.jshash2countfile, 'int')

	def reInit(self, jslistfile = 'bjslist.log', jshashfile = 'bjshashs', jshash2countfile = 'bjshash2count'):
		# for benign urls
		self.jslistfile = self.basedir + jslistfile
		self.jshashfile = self.basedir + jshashfile
		self.jshash2countfile = self.basedir + jshash2countfile
		self.jslist = utils.DictFile().file2dict(self.jslistfile, 'None').keys()
		self.jslist = [item.split('/')[-1] for item in self.jslist]
		self.jshashs = utils.DictFile().file2dict(self.jshashfile, 'str')
		self.jshash2count = utils.DictFile().file2dict(self.jshash2countfile, 'int')

	def saveList(self, content, filename):
		print "Saving list to {0}......".format(filename)
		outf = open(filename, 'w')
		for sublist in content:
			outf.write('\n'.join(sublist) + '\n')
		outf.close()

	def cpSampleList(self, sampleList, indir, outdir):
		print "Copying sample list from {0} to {1}".format(indir, outdir)
		utils.checkandaskdir(outdir)
		inf = open(sampleList, 'r')
		for line in inf:
			if line == '\n':
				continue
			dirname = line[:-1]
			# print 'Processing {0}.............'.format(dirname)
			cpcmd = 'cp -r ' + indir + dirname + ' ' + outdir
			commands.getoutput(cpcmd)

	def isFailedUrl(self, dirname):
		dirname = '{0}/{1}/'.format(self.groundTruthPath, dirname)
		# reloadUrl, failedUrl
		grepcmd = "grep 'failedUrl' {0}*".format(dirname)
		if not commands.getoutput(grepcmd):
			return False 
		return True

	def isValid(self, dirname):
		urlhash = dirname.split('.')[-1]
		if urlhash not in self.jshashs:
			return False
		return False if ( (self.jshash2count[self.jshashs[urlhash]] > 1) or \
				self.isFailedUrl(dirname) ) else True

	def getFilenameForClass(self, classname, count):
		# When selecting the benign examples, reInit
		if classname == 'B1':
			self.reInit()
		# Randomly select count urls
		self.groundTruthPath = '../../data/JSAdd2Tree/'
		filenames = utils.lsresult("-d {0}/{1}.*".format(self.groundTruthPath, classname))
		random.shuffle(filenames)
		# Handle invalid urls here
		progress = 0
		result = list()
		for filename in filenames:
			if progress >= count:
				return result
			if self.isValid(filename):
				result.append(filename)
				progress = progress + 1
		return result

	def pickGroundTruth(self, malClassList, numPerClass=1000, includeBenign=True):
		# The benign class is included by default
		# The format for malClassList is:
		# type,label,totalCount
		#
		# @type, name of a malicious class
		# @label, for example, M123, a string to identify type
		# @totalCount, availble number of urls for this class
		allClasses = list()
		allFilenames = list()
		toUseCount = dict()
		for line in open(malClassList, 'r'):
			line = line.split(',')
			totalCount = int(line[2])
			toUseCount[line[0]] = min(totalCount, numPerClass)
			allClasses.append(line[0])

		# Add benign classes
		if includeBenign:
			# Assuming there are at least numPerClass Benign ones there
			label = 'B1'
			allClasses.append(label)
			toUseCount[label] = numPerClass
		# Traverse classes and select samples
		for label in allClasses:
			# Get toUseCount for each class.
			# Handle invalid urls here.
			result = self.getFilenameForClass(label, toUseCount[label])
			allFilenames.append(result)
			print "Selected {0} urls for class {1}".format(len(result), label)
		# Save the sample list and copy to specified directory
		sampleList = 'samplelist'
		sampleList = self.basedir + sampleList
		self.saveList(allFilenames, sampleList)
		self.cpSampleList(sampleList, '/space/webcrawl-malware/data/JSAdd2Tree/', '/space/webcrawl-malware/data/MultipleClass/')


if __name__ == '__main__':
	gt = GroundTruth()
	# gt = GroundTruth('/home/ruian/chromium/webcrawl-malware/data/TempData/LOGs/')
	gt.pickGroundTruth(malClassList="/space/webcrawl-malware/data/TempData/URLs/top20.malurl.stats", numPerClass=100)
	# gt.pickGroundTruth(malClassList="/home/ruian/chromium/webcrawl-malware/data/TempData/URLs/top20.malurl.stats")

